'''
数据集分割, 训练、验证比例为8:2
'''
import os
import PIL.Image as Image
import numpy as np
import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *

# root = 'G:/USVInland/Stereo Matching/Low_Res_640_320' # Windows
# root = '/media/ubuntu/e/zhouyiqing/USVInland/Stereo Matching/Low_Res_640_320' # Ubuntu

# USVInland数据集
class USVInland(data.Dataset):
    def __init__(self, root, transform, mode):
        self.root = root
        self.transform = transform
        self.mode = mode

        # 拼接左右图像文件夹路径
        if seg: # 分割水体
            left_path = os.path.join(root, 'Left_Img_Rectified_Seg_rm')
            right_path = os.path.join(root, 'Right_Img_Rectified_Seg_rm')
            disp_path = os.path.join(root, 'Disp_Map_Seg_rm')
        else: # 未分割
            left_path = os.path.join(root, 'Left_Img_Rectified')
            right_path = os.path.join(root, 'Right_Img_Rectified')
            disp_path = os.path.join(root, 'Disp_Map')

        vali_range = [0, 1, 4, 12, 21, 27, 29, 31, 43, 47,
              49, 57, 62, 67, 69, 76, 80, 90, 95, 100,
              110, 113, 119, 123, 125, 130, 135, 138, 146, 148,
              158, 159, 163, 167, 170, 176] # validation set indices (36 images in total)

        # 对应图片在文件夹的范围
        if mode == 'train':
            imgs_range = [i for i in range(180) if i not in vali_range]
        elif mode == 'validate':
            imgs_range = vali_range

        left_imgs = list()
        right_imgs = list()
        disp_imgs = list()

        # 获取所有图像，并加入list
        all_left_imgs = list(os.path.join(left_path, img) for img in os.listdir(left_path))
        all_right_imgs = list(os.path.join(right_path, img) for img in os.listdir(right_path))
        all_disp_imgs = list(os.path.join(disp_path, img) for img in os.listdir(disp_path))

        # 按照图片范围加入list
        for i in imgs_range:
            left_imgs.append(all_left_imgs[i])
            right_imgs.append(all_right_imgs[i])
            disp_imgs.append(all_disp_imgs[i])

        self.left_imgs = left_imgs
        self.right_imgs = right_imgs
        self.disp_imgs = disp_imgs

    def __getitem__(self, index):
        # 读取RGB图像
        left_img = Image.open(self.left_imgs[index]).convert('RGB')
        right_img = Image.open(self.right_imgs[index]).convert('RGB')
        if system == 'Ubuntu': img_name = self.left_imgs[index].split('/')[-1]
        else: img_name = self.left_imgs[index].split('\\')[-1]

        # 读取视差灰度图像
        disp_img = Image.open(self.disp_imgs[index])

        # 图像预处理
        if self.transform:
            seed = torch.random.seed()
            torch.random.manual_seed(seed)
            left_img = self.transform(left_img)
            torch.random.manual_seed(seed)
            right_img = self.transform(right_img)

            disp_transform = transforms.Compose([transforms.RandomCrop((height, width))]) # 随机裁剪
            torch.random.manual_seed(seed)
            disp_img = np.array(disp_transform(disp_img), dtype=np.float32) / 255 * 50
            if under_sampling:
                disp_img = disp_img // 2

            # 张量转图片显示
            # transforms.ToPILImage()(left_img).convert('RGB').show()
            # transforms.ToPILImage()(right_img).convert('RGB').show()
            # transforms.ToPILImage()(disp_img).convert('L').show()

            # 数据保存到字典
            data = {'left': left_img, 'right': right_img, 'disp': disp_img, 'name': img_name}

        return data

    def __len__(self):
        return len(self.left_imgs)
